package graph;

import union_find_set.UnionFindSet;

import java.util.Comparator;
import java.util.HashSet;
import java.util.PriorityQueue;
import java.util.Set;

public class Kruskal_23111i5 {
    public static Set<Edge<Integer>> kruskal(Graph<Integer> graph) {
        PriorityQueue<Edge<Integer>> pq = new PriorityQueue<>(Comparator.comparingInt(key -> key.weight));
        // M 边数
        // O(M*logM)
        for (Edge<Integer> edge : graph.edges) {
            pq.offer(edge);
        }
        UnionFindSet<Node<Integer>> ufs = new UnionFindSet<>(graph.nodeMap.values());
        // 结果集，存储最小生成树的所有边
        Set<Edge<Integer>> result = new HashSet<>();
        // 存储已经冻结的节点
        Set<Node<Integer>> nodeSet = new HashSet<>(graph.nodeMap.size());
        // O(M*logM)
        // 小根堆不为空，且还没有找到所有节点，就继续循环
        while (!pq.isEmpty() && nodeSet.size() != graph.nodeMap.size()) {
            // 每次弹出最小的边
            Edge<Integer> minEdge = pq.poll();
            if (!ufs.isSameSet(minEdge.from, minEdge.to)) {
                ufs.union(minEdge.from, minEdge.to);
                result.add(minEdge);
                nodeSet.add(minEdge.from);
                nodeSet.add(minEdge.to);
            }
        }
        return result;
    }
}
